package com.feidee.fd.sml.algorithm.component.ml.classification

import com.feidee.fd.sml.algorithm.component.ml.{AbstractMLComponent, MLParam}
import com.feidee.fd.sml.algorithm.forecast.StageFinder
import com.feidee.fd.sml.algorithm.util.Constants
import com.feidee.fdspark.transformer.{ColEliminator, MetaStorage, ModelType}
import org.apache.spark.SparkException
import org.apache.spark.ml.evaluation.{BinaryClassificationEvaluator, MulticlassClassificationEvaluator}
import org.apache.spark.ml.feature.{IndexToString, StringIndexer, StringIndexerModel}
import org.apache.spark.ml.{Pipeline, PipelineModel}
import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions.udf
import org.apache.spark.sql.types.StringType

/**
  * @Author songhaicheng
  * @Date 2019/3/26 16:36
  * @Description
  * @Reviewer
  */
abstract class AbstractClassificationComponent[A <: MLParam](implicit m:Manifest[A])
  extends AbstractMLComponent[A] with Serializable {

  /**
    * 训练分类算法模型
    * @param param 算法运行参数
    * @param data  输入（训练用）数据
    * @return      算法模型
    */
  override def train(param: A, data: DataFrame): PipelineModel = {
    // 检查特征列和标签列是否存在
    val fieldNames = data.schema.fieldNames
    if (!fieldNames.contains(param.featuresCol)) {
      throw new SparkException(s"特征列 ${param.featuresCol} 不存在训练数据列 ${fieldNames.mkString("[", ", ", "]")} 中")
    }
    if (!fieldNames.contains(param.labelCol)) {
      throw new SparkException(s"标签列 ${param.labelCol} 不存在训练数据列 ${fieldNames.mkString("[", ", ", "]")} 中")
    }


    // 存储训练参数元信息
    val meta = new MetaStorage()
      .setModelType(ModelType.Algorithm_Classification_Prob)
      .setParameters(param.toMap)
      .setFields(data.schema.fieldNames)

    // 将原参数设置的 labelCol 和 predictionCol 改成中间临时名
    // 以实现标签索引、反索引，以及方便开发算法组件时直接调用 param 进行参数设置
    val existingColNames = data.schema.fieldNames
    val oriLabelCol = param.labelCol
    val oriPredictionCol = param.predictionCol
    val renamedLabelCol = tool.renameDuplicatedColName(s"${param.labelCol}_rename", existingColNames)
    val renamedPredictionCol = tool.renameDuplicatedColName(s"${param.predictionCol}_rename", existingColNames)
    param.labelCol = renamedLabelCol
    param.predictionCol = renamedPredictionCol

    // 下面两个步骤方便训练模型和还原真实的标签
    val indexer = new StringIndexer()
      .setHandleInvalid("error")         // 设为 "error"，防止例如 LinearSVC 组件错误识别标签类别数量（expect 2 but has 3）
      .setInputCol(oriLabelCol)
      .setOutputCol(param.labelCol)
      .fit(data)

    val converter = new IndexToString()
      .setInputCol(param.predictionCol)
      .setOutputCol(oriPredictionCol)
      .setLabels(indexer.labels)

    // 最后返回的 DataFrame 需要把中间过程辅助列去掉，并将结果列还原成参数设置的样子
    val colEliminator = new ColEliminator()
      .setDrops(Array[String](param.labelCol, param.predictionCol))

    // 执行流程：保存训练参数元数据 》 标签索引化 》 具体算法组件 》 还原标签值 》 去掉辅助列
    val pipeline = new Pipeline()
      .setStages(Array(meta, indexer, setUp(param), converter, colEliminator))
    val model = pipeline.fit(data)

    // 还原 labelCol 和 predictionCol 为原来设置的值，防止后续需要借助读取参数实现的方法（如计算指标）因拿不到正确参数而报错
    param.labelCol = oriLabelCol
    param.predictionCol = oriPredictionCol

    model
  }

  /**
    * 根据参数，计算分类参数指标
    * @param param    模型训练时的参数
    * @param labels   分类模型
    * @param data     评估用的 DataFrame
    * @return         指标值
    */
  override def calculateMetrics(param: A, model: PipelineModel, data: DataFrame): Map[String, Double] = {
    var metrics = Map[String, Double]()

    // 从生成的分类模型里提取出标签信息
    val labels = new StageFinder[StringIndexerModel](model).findWithOrder().get.labels
    // 把预测结果转成 double 类型
    val getRawLabel = udf{ label: String => {
      labels.indexOf(label).toDouble
    }}
    val rawData = data
      .withColumn(param.predictionCol, getRawLabel(data.col(param.predictionCol).cast(StringType)))
      .withColumn(param.labelCol, getRawLabel(data.col(param.labelCol).cast(StringType)))

    val bEvaluator = new BinaryClassificationEvaluator()
      .setLabelCol(param.labelCol)
    val mEvaluator = new MulticlassClassificationEvaluator()
      .setPredictionCol(param.predictionCol)
      .setLabelCol(param.labelCol)

    for (i <- param.metrics.indices) {
      param.metrics(i).toLowerCase match {
        case "accuracy" | "precision" | "recall" | "f1" =>
          mEvaluator.setMetricName(Constants.classificationMetricNames(param.metrics(i).toLowerCase))
          metrics += (param.metrics(i) -> mEvaluator.evaluate(rawData))
        case "auc" | "pr" =>
          // 二分类计算指标需要传入 probability 列
          require(param.toMap.get("probabilityCol").nonEmpty,
            "unsupported binary metric [auc and pr] since lacking of probabilityCol")
          val probabilityCol = param.toMap("probabilityCol").toString
          bEvaluator
            .setRawPredictionCol(probabilityCol)
            .setMetricName(Constants.classificationMetricNames(param.metrics(i).toLowerCase))
          metrics += (param.metrics(i) -> bEvaluator.evaluate(rawData))
        case _ => logInfo(s"UNKNOWN METRIC ${param.metrics(i)}")
      }
    }

    metrics
  }

}
